#ifndef __UNIONFINDSET_HPP__
#define __UNIONFINDSET_HPP__

#include <iostream>
#include <vector>
#include <unordered_map>
#include <iostream>
using namespace std;


namespace clx_datastructure {
    template<class V, int DFLT_VAL = -1>
    class UnionFindSet {
        public:
            UnionFindSet(vector<V>& vertexs) : ufs(vertexs.size(), DFLT_VAL) {
                for (int i = 0; i < vertexs.size(); i++) {
                    vIndexMap[vertexs[i]] = i;
                }
            }

            int getIndex(const V& v) {
                if (vIndexMap.find(v) == vIndexMap.end()) {
                    return -1;
                }
                return vIndexMap[v];
            }

            int findRoot(const V& v) {
                int num = getIndex(v);
                while (ufs[num] >= 0) {
                    num = ufs[num];
                }
                return num;
            }

            int findRoot_v2(const V& v) {
                int num = getIndex(v);
                int tmp = num;
                while (ufs[tmp] > 0) {
                    tmp = ufs[tmp];
                }

                while (ufs[num] > 0) {
                    int parent = ufs[num];
                    ufs[num] = tmp;
                    num = parent;
                }
            }

            bool inSameSet(const V& v1, const V& v2) {
                return findRoot(v1) == findRoot(v2);
            }

            bool unionSet(const V& v1, const V& v2) {
                int root1 = findRoot(v1), root2 = findRoot(v2);
                if (root1 == root2) return false;

                if (root2 < root1) 
                    swap(root1, root2);

                ufs[root1] += ufs[root2];
                ufs[root2] = root1;
                return true;
            }

            int getNum() {
                int count = 0;
                for (int i = 0; i < ufs.size(); i++) {
                    if (ufs[i] < 0) count++;
                }
                return count;
            }

            void ufs_print() {
                for (int i = 0; i < ufs.size(); i++) 
                    cout << ufs[i] << " ";
                cout << endl;
            }

        private:
            vector<int> ufs;
            unordered_map<V, int> vIndexMap;
    };
}
#endif